#!/usr/bin/env python
# coding=utf-8
# __author__ = 'Yunchao Ling'

from pymongo import MongoClient
from datetime import datetime
import fire


# 连接MongoDB
def InitMongoDB(MONGODB_HOST, MONGODB_PORT, MONGODB_DB_NAME, MONGODB_COLLECTION_NAME):
    client = MongoClient(MONGODB_HOST, MONGODB_PORT)
    db = client[MONGODB_DB_NAME]
    collection = db[MONGODB_COLLECTION_NAME]
    return client, collection


# 关闭MongoDB
def CloseMongoDB(MongoDB_Connection):
    MongoDB_Connection.close()


def TestDownloader():
    downloader_list = {}
    infile_downloader = open("D:\public_data_download\previous_list\downloader.txt", "r")
    for line in infile_downloader:
        line = line.rstrip()
        splitline = line.split("\t")
        downloader_list[splitline[0] + ".runid"] = splitline[1]
    infile_downloader.close()

    no_run = 0
    no_downloader = 0
    same_downloader = 0
    diff_downloader = 0

    client, collection = InitMongoDB("10.10.116.28", 27017, "PD_Data", "rawdata")

    # outfile = open("D:/TestDownloader.txt", "w")

    for filename in downloader_list.keys():
        infile = open("D:\public_data_download\previous_list\\" + filename, "r", encoding="utf-8")
        for line in infile:
            line = line.rstrip()
            result = collection.find_one({"RunID": line}, {"downloader": 1})
            if result == None:
                no_run += 1
                # print(line)
            elif "downloader" not in result:
                no_downloader += 1
            elif result["downloader"] != downloader_list[filename]:
                diff_downloader += 1
                # print("[Diff Downloader]file downloader=" + downloader_list[filename] + ";db downloader=" + result[
                #     "downloader"])
                # outfile.write(
                #     "[Diff Downloader]file downloader=" + downloader_list[filename] + ";db downloader=" + result[
                #         "downloader"] + "\n")
                # outfile.flush()
            else:
                same_downloader += 1
        infile.close()

    print("[no run]" + str(no_run))
    print("[no_downloader]" + str(no_downloader))
    print("[same_downloader]" + str(same_downloader))
    print("[diff_downloader]" + str(diff_downloader))

    # outfile.write("[no run]" + str(no_run))
    # outfile.write("[no_downloader]" + str(no_downloader))
    # outfile.write("[same_downloader]" + str(same_downloader))
    # outfile.write("[diff_downloader]" + str(diff_downloader))
    #
    # outfile.flush()
    # outfile.close()

    CloseMongoDB(client)


def FullfillJob():
    downloader_list = {}
    infile_downloader = open("downloader.txt", "r")
    for line in infile_downloader:
        line = line.rstrip()
        splitline = line.split("\t")
        downloader_list[splitline[0] + ".runid"] = splitline[1]
    infile_downloader.close()

    no_run = 0
    no_downloader = 0
    same_downloader = 0
    diff_downloader = 0

    client, collection = InitMongoDB("localhost", 27017, "PD_Data", "rawdata")
    client2, collection2 = InitMongoDB("localhost", 27017, "SRA_metadata", "accessions")

    # outfile = open("D:/TestDownloader.txt", "w")

    for filename in downloader_list.keys():
        infile = open("previous_list/" + filename, "r", encoding="utf-8")
        for line in infile:
            line = line.rstrip()
            result = collection.find_one({"RunID": line}, {"downloader": 1})
            if result == None:
                no_run += 1
                insert_record = {}
                insert_record["RunID"] = line
                result2 = collection2.find_one({"Accession": line}, {"BioProject": 1, "Study": 1})
                if result2 != None:
                    if "BioProject" in result2:
                        insert_record["BioprojectID"] = result2["BioProject"]
                    if "Study" in result2:
                        insert_record["Study"] = result2["Study"]
                insert_record["source"] = "SRA"
                insert_record[
                    "data_download_url"] = "anonftp@ftp.ncbi.nlm.nih.gov:/sra/sra-instant/reads/ByRun/sra/" + line[
                                                                                                              :3] + "/" + line[
                                                                                                                          :6] + "/" + line + "/" + line + ".sra"
                insert_record["download_count"] = "1"
                insert_record["Date_to_Yao"] = datetime.today()
                insert_record["download_way"] = "external"
                insert_record["downloader"] = downloader_list[filename]
                insert_record["JobID"] = filename.split(".")[0]
                collection.insert_one(insert_record)
            elif "downloader" not in result:
                no_downloader += 1
                collection.update_one({"_id": result["_id"]}, {
                    "$set": {"downloader": downloader_list[filename], "JobID": filename.split(".")[0]}})
            elif result["downloader"] != downloader_list[filename]:
                diff_downloader += 1
                collection.update_one({"_id": result["_id"]}, {
                    "$set": {"downloader": downloader_list[filename], "JobID": filename.split(".")[0]}})
                # print("[Diff Downloader]file downloader=" + downloader_list[filename] + ";db downloader=" + result[
                #     "downloader"])
                # outfile.write(
                #     "[Diff Downloader]file downloader=" + downloader_list[filename] + ";db downloader=" + result[
                #         "downloader"] + "\n")
                # outfile.flush()
            else:
                same_downloader += 1
                collection.update_one({"_id": result["_id"]}, {
                    "$set": {"JobID": filename.split(".")[0]}})
        infile.close()

    print("[no run]" + str(no_run))
    print("[no_downloader]" + str(no_downloader))
    print("[same_downloader]" + str(same_downloader))
    print("[diff_downloader]" + str(diff_downloader))

    # outfile.write("[no run]" + str(no_run))
    # outfile.write("[no_downloader]" + str(no_downloader))
    # outfile.write("[same_downloader]" + str(same_downloader))
    # outfile.write("[diff_downloader]" + str(diff_downloader))
    #
    # outfile.flush()
    # outfile.close()

    CloseMongoDB(client)
    CloseMongoDB(client2)


def TestEBI():
    client, collection = InitMongoDB("10.10.116.28", 27017, "PD_Data", "rawdata")
    client2, collection2 = InitMongoDB("10.10.116.28", 27017, "SRA_metadata", "accessions")

    for result in collection.find({"source": "EBI"}, {"RunID": 1}):
        runid = result['RunID']
        result2 = collection2.find_one({"Accession": runid}, {"BioProject": 1, "Study": 1})
        if result2 != None:
            # printline = runid + "|"
            # if "BioProject" in result2:
            #     printline += result2["BioProject"]
            # printline += "|"
            # if "Study" in result2:
            #     printline += result2["Study"]
            # print(printline)
            collection.update_one({"_id": result["_id"]},
                                  {"$set": {"Study": result2["Study"]}, "$unset": {"BioprojectID": 1}})
        else:
            print("No run")

    CloseMongoDB(client)
    CloseMongoDB(client2)


if __name__ == '__main__':
    # TestDownloader()
    fire.Fire()
